package com.cn.query;

import com.alibaba.fastjson.JSONObject;
import com.cn.client.SearchClient;
import com.google.common.collect.Maps;
import org.apache.http.HttpHost;
import org.elasticsearch.action.search.MultiSearchRequest;
import org.elasticsearch.action.search.MultiSearchResponse;
import org.elasticsearch.action.search.SearchRequest;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.client.*;
import org.elasticsearch.common.lucene.search.function.CombineFunction;
import org.elasticsearch.common.unit.TimeValue;
import org.elasticsearch.index.query.MatchAllQueryBuilder;
import org.elasticsearch.index.query.QueryBuilders;
import org.elasticsearch.index.query.TermQueryBuilder;
import org.elasticsearch.index.query.functionscore.FunctionScoreQueryBuilder;
import org.elasticsearch.index.query.functionscore.ScriptScoreFunctionBuilder;
import org.elasticsearch.script.Script;
import org.elasticsearch.search.SearchHit;
import org.elasticsearch.search.builder.SearchSourceBuilder;
import org.elasticsearch.search.sort.*;
import org.junit.Before;
import org.junit.Test;

import java.io.IOException;
import java.text.SimpleDateFormat;
import java.util.Collections;
import java.util.Map;

import static org.elasticsearch.script.Script.DEFAULT_SCRIPT_LANG;
import static org.elasticsearch.script.Script.DEFAULT_SCRIPT_TYPE;

public class ESSearchDemo {

  private static SimpleDateFormat simpleDateFormat = new SimpleDateFormat("yyyyMMdd");

  static RestHighLevelClient client;

  @Before
  public void before() {
    RestClientBuilder builder = RestClient.builder(
            new HttpHost("10.211.55.3", 9200, "http"));
    client = new RestHighLevelClient(builder);
  }

  @Test
  public void test11() {
    //查询所有
    SearchRequest searchRequest = new SearchRequest();
    SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
    searchSourceBuilder.query(QueryBuilders.matchAllQuery());
    searchRequest.source(searchSourceBuilder);
    SearchResponse search = null;
    try {
      search = client.search(searchRequest, RequestOptions.DEFAULT);
    } catch (IOException e) {
      e.printStackTrace();
    }
    println(search);
  }

  @Test
  public void test0() {
    try {
      SearchRequest searchRequest = new SearchRequest();
      SearchResponse search = client.search(searchRequest, RequestOptions.DEFAULT);
      println(search);
    } catch (IOException e) {
      e.printStackTrace();
    }
  }

  @Test
  public void nestedsortTest() throws IOException {

//query
    SearchSourceBuilder sourceBuilder = new SearchSourceBuilder();
    MatchAllQueryBuilder matchAllQueryBuilder = QueryBuilders.matchAllQuery();
    sourceBuilder.query(matchAllQueryBuilder);
//嵌套类型打分打分排序
    FieldSortBuilder fieldSortBuilder = SortBuilders.fieldSort("multirank.v")
            .order(SortOrder.DESC)
            .sortMode(SortMode.AVG)
            .setNestedSort(new NestedSortBuilder("multirank")
                    .setFilter(QueryBuilders.termQuery("multirank.k", 2))

            );

    sourceBuilder.sort(fieldSortBuilder);
    SearchRequest searchRequest = new SearchRequest("my_store");
    searchRequest.source(sourceBuilder);
    search(searchRequest);
  }

  @Test
  //LocationAndNameFunction
  public void test122() {
    try {

      Map<String, Object> params = Maps.newHashMap();
      params.put("_lat", 30.22);
      params.put("_lon", 120.1);
      params.put("_cityid", 10);
      params.put("_maxdistance", 300);
      params.put("_keyword", "keyword");
      Map<String, Object> params1 = Maps.newHashMap();
      params1.put("params", params);

      String script =
              "int keywordscore = doc['word'].value.equals(params._keyword) ? 10:0;" +
                      "int cityidscore = (doc['cityid'] == params._cityid) ? 10:0;" +
                      "double lat = doc['location'].lat ;" +
                      "double lon = doc['location'].lon ;" +
                      "double diffLon = Math.abs(lon - params._lon);" +
                      "diffLon = (diffLon > 180)?diffLon-= 360:diffLon;" +
                      "double distance = Math.sqrt(Math.pow(diffLon, 2) + Math.pow(lat - params._lat, 2)) * 110.0;" +
                      "double distancescore = (distance<0||distance > params._maxdistance)?0:10 * (1 - distance / params._maxdistance);" +
                      "return distancescore + keywordscore + cityidscore;";

      ScriptScoreFunctionBuilder scriptScoreFunctionBuilder = new ScriptScoreFunctionBuilder(
              new Script(DEFAULT_SCRIPT_TYPE, DEFAULT_SCRIPT_LANG, script
                      , Collections.emptyMap()
                      , params)
      );

//functionscore
      FunctionScoreQueryBuilder functionScoreQueryBuilder = QueryBuilders.functionScoreQuery(
              QueryBuilders.boolQuery()
                      .must(QueryBuilders.matchQuery("keyword", "酒店")).must(QueryBuilders.termQuery("type", "hotel"))
              , scriptScoreFunctionBuilder).boostMode(CombineFunction.SUM);
//query
      SearchSourceBuilder sourceBuilder = new SearchSourceBuilder();
      sourceBuilder.query(functionScoreQueryBuilder);
//打分排序
      sourceBuilder.sort(SortBuilders.scoreSort());

      SearchRequest searchRequest = new SearchRequest("globalsearch");
      searchRequest.source(sourceBuilder);

      search(searchRequest);
    } catch (Exception e) {
      e.printStackTrace();
    }
  }

  private void search(SearchRequest searchRequest) throws IOException {
    MultiSearchRequest multiSearchRequest = new MultiSearchRequest();
    multiSearchRequest.add(searchRequest);
    MultiSearchResponse multiSearchResponse = client.msearch(
            multiSearchRequest, RequestOptions.DEFAULT);
    System.out.println(multiSearchResponse);
    MultiSearchResponse.Item[] responses = multiSearchResponse.getResponses();
    TimeValue took = multiSearchResponse.getTook();
    System.out.println(took);
  }

  //ScoreAndDistanceFunction
  public static void test2() {
    try {
      MultiSearchRequest multiSearchRequest = new MultiSearchRequest();
      SearchSourceBuilder sourceBuilder = new SearchSourceBuilder();
      Map<String, Object> params = Maps.newHashMap();
      params.put("_lat", 30.22);
      params.put("_lon", 120.1);
      params.put("_maxdistance", 100);
      Map<String, Object> params1 = Maps.newHashMap();
      params1.put("params", params);
//排序 脚步
      String script =
              "double lat = doc['location'].lat ;" +
                      "double lon = doc['location'].lon ;" +
                      "double diffLon = Math.abs(lon - params._lon);" +
                      "diffLon = (diffLon > 180)?diffLon-= 360:diffLon;" +
                      "double distance = Math.sqrt(Math.pow(diffLon, 2) + Math.pow(lat - params._lat, 2)) * 110.0;" +
                      "double distancescore = (distance<0||distance > params._maxdistance)?0:10 * (1 - distance / params._maxdistance);" +
                      "return distancescore*0.6 +doc['rankscore'].value*0.4 ;";
      script = String.format(script, 30.87, 120.87);
      sourceBuilder.sort(SortBuilders.scriptSort(
              new Script(DEFAULT_SCRIPT_TYPE, DEFAULT_SCRIPT_LANG, script, Collections.emptyMap(), params)
              , ScriptSortBuilder.ScriptSortType.NUMBER)
              .order(SortOrder.DESC)
      );
//条件
      sourceBuilder.query(QueryBuilders.boolQuery().must(new TermQueryBuilder("type", "hotel")));

      SearchRequest searchRequest = new SearchRequest("globalsearch");
      searchRequest.source(sourceBuilder);
      multiSearchRequest.add(searchRequest);
      MultiSearchResponse multiSearchResponse = client.msearch(
              multiSearchRequest, RequestOptions.DEFAULT);
      System.out.println(multiSearchResponse);
      MultiSearchResponse.Item[] responses = multiSearchResponse.getResponses();
      TimeValue took = multiSearchResponse.getTook();
      System.out.println(took);
    } catch (Exception e) {
      e.printStackTrace();
    }
  }

  private void println(SearchResponse search) {
    if (search != null) {
      for (SearchHit hit : search.getHits()) {
        Map<String, Object> sourceAsMap = hit.getSourceAsMap();
        System.out.println("docid:" + hit.getId());
      }

    }
  }

  @Test
  public void singlesearch() {
    try {
      SearchRequest searchRequest = new SearchRequest();

      SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
      searchSourceBuilder.query(QueryBuilders.matchQuery("word", "可乐"));
      searchRequest.source(searchSourceBuilder);
      searchRequest.indices("my_*");
      SearchResponse search = SearchClient.client().search(searchRequest, RequestOptions.DEFAULT);
      JSONObject jsonObject = new JSONObject();
      System.out.println(jsonObject.toJSONString(search));
      for (SearchHit hit : search.getHits()) {

        Map<String, Object> sourceAsMap = hit.getSourceAsMap();
        System.out.println(jsonObject.toJSONString(sourceAsMap));
      }
    } catch (IOException e) {
      e.printStackTrace();
    }
  }

  @Test
  public void msearch() {

    MultiSearchRequest request = new MultiSearchRequest();
    SearchRequest firstSearchRequest = new SearchRequest();
    SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
    searchSourceBuilder.query(QueryBuilders.matchQuery("user", "kimchy"));
    firstSearchRequest.source(searchSourceBuilder);
    request.add(firstSearchRequest);
    SearchRequest secondSearchRequest = new SearchRequest();
    searchSourceBuilder = new SearchSourceBuilder();
    searchSourceBuilder.query(QueryBuilders.matchQuery("word", "可乐"));
    secondSearchRequest.source(searchSourceBuilder);
    request.add(secondSearchRequest);

    try {
      MultiSearchResponse msearch = SearchClient.client().msearch(request, RequestOptions.DEFAULT);
      for (int i = 0; i < msearch.getResponses().length; i++) {
        MultiSearchResponse.Item respons = msearch.getResponses()[i];

        JSONObject jsonObject = new JSONObject();

        for (SearchHit hit : respons.getResponse().getHits()) {
          Map<String, Object> sourceAsMap = hit.getSourceAsMap();
          System.out.println(jsonObject.toJSONString(sourceAsMap));
          String id = hit.getId();
          float score = hit.getScore();
        }
      }
    } catch (IOException e) {
      e.printStackTrace();
    }

  }


}
